#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *
from datetime import datetime, timezone

parser = argparse.ArgumentParser(description='Plot nonhyrostatic reduced sphere mountain wave test results')
parser.add_argument('-i', '--input', help='Input data file', default='ksp15_02.h0.nc')
parser.add_argument('-t', '--time-step', type=int, default=-1, help='Time step to plot')
parser.add_argument('-o', '--output', help='Output figure')
parser.add_argument('-v', '--var', help='Variable to plot', default='w')
parser.add_argument('-l', '--loop', help='Loop the plot', action='store_true')
args = parser.parse_args()

fig = plt.figure(figsize=(15, 10))
proj = ccrs.PlateCarree(central_longitude=180)
zlev = np.linspace(0.1, 19.9, 101) * units.km

f = addfile(args.input).isel(time=0)
pt0_1 = vinterp_z(f.gz.isel(lat=180) * units('m**2 / s**2') / g, 300 * units.K * (1e5 / f.ph.isel(lat=180))**(rd/cp), zlev)
pt0_2 = vinterp_z(f.gz * units('m**2 / s**2') / g, 300 * units.K * (1e5 / f.ph)**(rd/cp), [8] * units.km)
f.close()

levels_1 = np.linspace(-1, 1, 21)
colors_1 = ListedColormap([
	[0.15, 0.18, 0.36],
	[0.16, 0.27, 0.51],
	[0.23, 0.39, 0.66],
	[0.35, 0.55, 0.78],
	[0.49, 0.70, 0.87],
	[0.63, 0.82, 0.94],
	[0.76, 0.89, 0.97],
	[0.87, 0.94, 0.97],
	[0.95, 0.97, 1.00],
	[1.00, 1.00, 1.00],
	[1.00, 1.00, 1.00],
	[0.99, 0.97, 0.88],
	[0.98, 0.93, 0.67],
	[0.96, 0.82, 0.47],
	[0.94, 0.68, 0.30],
	[0.90, 0.53, 0.26],
	[0.86, 0.38, 0.23],
	[0.79, 0.25, 0.20],
	[0.66, 0.19, 0.18],
	[0.53, 0.15, 0.14],
])

levels_2 = np.linspace(-0.6, 0.6, 13)
colors_2 = ListedColormap([
	[0.15, 0.18, 0.36],
	[0.20, 0.35, 0.64],
	[0.42, 0.64, 0.82],
	[0.67, 0.84, 0.95],
	[0.86, 0.93, 0.97],
	[1.00, 1.00, 1.00],
	[1.00, 1.00, 1.00],
	[0.98, 0.93, 0.64],
	[0.94, 0.70, 0.31],
	[0.89, 0.45, 0.24],
	[0.76, 0.21, 0.20],
	[0.53, 0.15, 0.14],
])

def plot(data_path, fig):
	f = addfile(args.input).isel(time=args.time_step)

	time = parse_datetime(f.time)
	fig.suptitle(f'Time: {time.strftime("%H:%M")}')

	pt = vinterp_z(f.gz.isel(lat=180) * units('m**2 / s**2') / g, f.pt.isel(lat=180), zlev) - pt0_1
	pt.attrs['long_name'] = 'Potential temperature perturbation'
	ax = fig.add_subplot(2, 2, 1)
	plot_contour_lon(ax, pt, cmap=colors_1, levels=levels_1, ticks=levels_1, with_contour=True, contour_levels=levels_1[np.abs(levels_1)>=0.01])
	ax.set_xlim(165, 285)

	w = vinterp_z(f.gz_lev.isel(lat=180) * units('m**2 / s**2') / g, f.w_lev.isel(lat=180), zlev)
	ax = fig.add_subplot(2, 2, 3)
	plot_contour_lon(ax, w, cmap=colors_1, levels=levels_1, ticks=levels_1, with_contour=True, contour_levels=levels_1[np.abs(levels_1)>=0.01])
	ax.set_xlim(165, 285)

	pt = vinterp_z(f.gz * units('m**2 / s**2') / g, f.pt, [8] * units.km) - pt0_2
	pt.attrs['long_name'] = 'Potential temperature perturbation'
	ax = fig.add_subplot(2, 2, 2, projection=proj)
	plot_contour_map(ax, pt[0,:,:], cmap=colors_2, levels=levels_2, ticks=levels_2, with_contour=True, contour_levels=levels_2[np.abs(levels_2)>=0.01])
	ax.set_extent([165, 285, -45, 45], crs=ccrs.PlateCarree())

	w = vinterp_z(f.gz_lev * units('m**2 / s**2') / g, f.w_lev, [8] * units.km)
	ax = fig.add_subplot(2, 2, 4, projection=proj)
	plot_contour_map(ax, w[0,:,:], cmap=colors_2, levels=levels_2, ticks=levels_2, with_contour=True, contour_levels=levels_2[np.abs(levels_2)>=0.01])
	ax.set_extent([165, 285, -45, 45], crs=ccrs.PlateCarree())

	f.close()

plot(args.input, fig)
if args.output:
	plt.savefig(args.output, dpi=300)
elif args.loop:
	loop_plot(args.input, fig, plot)
else:
	plt.tight_layout()
	plt.show()
